""" Integro differential equation a*u_xx + b* u_x + c*u + l(conv(u, p) - u) """

import itertools
from numpy import floor, empty, sum as npsum
from numpy import asarray, zeros, meshgrid
from numpy.linalg import solve

from SecondOrderStatic1D import SecondOrderStatic1D


class PIDEStatic1D(SecondOrderStatic1D):

    """ Solve the boundary value problem. """

    __slots__ = ['m_jumpma', 'm_pide_bop', 'm_kernel', 'm_lambda']


    def __init__(self, nsegs, xleft, xright):
        """ Instantiate an object of partial integro differential type."""

        SecondOrderStatic1D.__init__(self, nsegs, xleft, xright)
        (self.m_lambda, self.m_kernel) = (None, None)
        (self.m_pide_bop, self.m_jumpma) = (None, None)

    def __assemble_jump(self):
        """
        Info: Use quadrature formulas to assemble jump matrix.
        Output: jump matrix and boundary coefficients.
        """

        # get shape function and quad points.

        quad_formula = self.get_quads()
        quadp = quad_formula.get_qpoints()

        # evaluate at quad points.

        sh_funct = self.get_sfunc()
        shfu_at_quadp = sh_funct.values(quadp)

        # Multiply pointwise shape functions at quad.

        phiijs = []
        lc_dof = sh_funct.dof()
        for (i, j) in itertools.product(range(0, lc_dof), range(0,
                lc_dof)):
            (phii, phij) = meshgrid(shfu_at_quadp[i], shfu_at_quadp[j])
            phiijs.append((phii * phij).flatten())

        # Evaluate kernel[:, 0] and kernel[0, :]

        grid = self.get_mesh()
        delta_x = grid.getDx()[0]
        num_cols = grid.get_nsegs() * (lc_dof - 1) + 1

        # assemble jump matrix.

        jump = zeros((num_cols, num_cols), 'float64')

        for (row, col) in itertools.product(range(0, num_cols),
                range(0, num_cols)):
            jump[row, col] = self.__jumpij(row, col, num_cols - 1,
                    delta_x, phiijs)

        toe_coeff = zeros((2 * num_cols - 1, 1))
        toe_coeff[:num_cols] = jump[:, 0][::-1].reshape((num_cols, 1))
        toe_coeff[num_cols - 1:] = jump[0, :].reshape((num_cols, 1))

        return (self.m_lambda * jump[1:-1, 1:-1], self.m_lambda * toe_coeff)

       # TODO improve on this matter.
        # fill upper diagonal term

        for col_id in range(0, num_cols):
            t_upper = toe_coeff[num_cols - 1 + col_id]
            for indx in range(0, num_cols - col_id):
                jump[indx, col_id + indx] = t_upper

        # fill lower diagonal terms

        for row_id in range(1, num_cols):
            t_lower = toe_coeff[num_cols - row_id - 1]
            for indx in range(0, num_cols - row_id):
                jump[row_id + indx, indx] = t_lower

        return (self.m_lambda * jump[1:-1, 1:-1], self.m_lambda * toe_coeff)

    def __eval_kernel(self, points):
        """ expect (n, 2) vector. """

        return self.m_kernel(points)


    def __jumpij(self, row, col, stop, dx, phiijs):
        """ computest (row, col) coefficient.

            phiijs ~ list of (n, 1) vectors.
                     it holds the upper diagonal matrix
                     \phi_i\phi_j(qpoints)
        """

        if 0. == self.m_lambda:
            return 0.

        quad_formula = self.get_quads()
        quadp = dx * quad_formula.get_qpoints() / 2.
        weights = quad_formula.get_weights()

        (x_weig, y_weig) = meshgrid(weights, weights)
        xy_weig = x_weig.flatten() * y_weig.flatten()

        # all possible kernel quad points

        (xcol, ycol) = dx * floor(asarray([row, col]) / 2.)
        (xcoor, ycoor) = meshgrid(xcol + quadp, ycol + quadp)
        (xcoor_1, ycoor_1) = (xcoor - dx, ycoor - dx)

        jump_int = 0
        kernel_00 = self.m_kernel((xcoor - ycoor).flatten())
        if 0 == row % 2 and 0 == col % 2:
            if row < stop and col < stop:
                jump_int = npsum(kernel_00 * phiijs[0] * xy_weig)

            if 0 < row and col < stop:
                kernel_1_0 = self.m_kernel((xcoor_1 - ycoor).flatten())
                jump_int += npsum(kernel_1_0 * phiijs[6] * xy_weig)

            if row < stop and 0 < col:
                kernel_0_1 = self.m_kernel((xcoor - ycoor_1).flatten())
                jump_int += npsum(kernel_0_1 * phiijs[2] * xy_weig)

            if 0 < row and 0 < col:
                kernel_1_1 = self.m_kernel((xcoor_1
                        - ycoor_1).flatten())
                jump_int += npsum(kernel_1_1 * phiijs[8] * xy_weig)
        elif 1 == row % 2 and 0 == col % 2:
            if col < stop:
                jump_int = npsum(kernel_00 * phiijs[3] * xy_weig)
            if 0 < col:
                kernel_0_1 = self.m_kernel((xcoor - ycoor_1).flatten())
                jump_int += npsum(kernel_0_1 * phiijs[5] * xy_weig)
        elif 0 == row % 2 and 1 == col % 2:
            if row < stop:
                jump_int = npsum(kernel_00 * phiijs[1] * xy_weig)
            if 0 < row:
                kernel_1_0 = self.m_kernel((xcoor_1 - ycoor).flatten())
                jump_int += npsum(kernel_1_0 * phiijs[7] * xy_weig)
        else:
            jump_int = npsum(kernel_00 * phiijs[4] * xy_weig)

        return dx * dx * jump_int / 4.  # this basis has support 4.

    def get_jump(self):
        """ returns jump matrix in full format. """

        if self.m_jumpma is None:
            (self.m_jumpma, self.m_pide_bop) = self.__assemble_jump()
        return (self.m_jumpma, self.m_pide_bop)

    def get_lambda(self):
        """ return jump intensity. """

        return self.m_lambda

    def get_pde(self):  # PDE contribution
        """ pde matrix is sparse format. """

        return self.get_stiff_dformat()  # PDE contribution

    def get_mass(self):  # Time step discretization.
        """ mass matrix is sparse format. """

        return self.get_mass_dformat()

    def reset(self):
        """
            Free memory and set local variables to None.
        """

        if self.m_lambda is not None:
            del self.m_lambda
            self.m_lambda = None

        if self.m_kernel is not None:
            del self.m_kernel
            self.m_kernel = None

        if self.m_jumpma is not None:
            del self.m_jumpma
            self.m_jumpma = None

    def set_pide_cos(self, diff, conv, reac, lamb):
        """ diffusion, convection, reaction,  jump  """

        super(PIDEStatic1D, self).set_cos(diff, conv, reac - lamb)
        self.m_lambda = lamb

    def set_cos(self, diff, conv, reac, d_is_real=True,
            c_is_real=True, r_is_real=True):
        """ make sure that the pde interface does not get called when
            solving a pide equation. """

        raise NotImplementedError, \
            'use set_pide_cos to set PIDE params.'

    def set_kernel(self, kernel):
        """ user define kernel. """

        self.m_kernel = kernel

    def solve(self):
        """ solve the boundary value problem """

        # Jump matrix and boundary operator.

        (pide_mat, pide_bop) = self.__assemble_jump()

        # Add PDE contribution.

        stiff_diag = self.get_pde()  # PDE contribution

        for (posi, val) in enumerate(stiff_diag[1:-1, 0]):
            pide_mat[posi, posi] += val

        for (posi, val) in enumerate(stiff_diag[2:-1, 1]):
            pide_mat[posi, posi + 1] += val

        for (posi, val) in enumerate(stiff_diag[1:-2, 2]):
            pide_mat[posi + 1, posi] += val

        if 3 == self.get_sfunc().dof():
            for (posi, val) in enumerate(stiff_diag[3:-1, 3]):
                pide_mat[posi, posi + 2] += val

            for (posi, val) in enumerate(stiff_diag[1:-3, 4]):
                pide_mat[posi + 2, posi] += val

        (uleft, uright) = self.get_bd_conditions()
        rhs = self.get_load_terms()

        # Complement RHS with jump contribution.

        left_jump = pide_bop[:rhs.size]
        right_jump = pide_bop[-rhs.size:]
        rhs -= (uleft * left_jump + uright * right_jump)[::-1]

        # Complement RHS with PDE contribution.

        if 3 == self.get_sfunc().dof():
            rhs[2] -= uleft * stiff_diag[0, 4]
            rhs[-3] -= uright * stiff_diag[-1, 3]

        rhs[1] -= uleft * stiff_diag[0, 2]
        rhs[-2] -= uright * stiff_diag[-1, 1]

        load = rhs[1:-1]

        # Solve the system of equations.

        solu = empty((rhs.size, 1))
        solu[1:-1] = solve(pide_mat, load).reshape((load.size, 1))
        (solu[0], solu[-1]) = (uleft, uright)
        return solu


